import numpy as np
import pyrender
import trimesh
import math
import cv2
from geometry_msgs.msg import Pose, PoseStamped, PoseArray, Point, Quaternion
from visualization_msgs.msg import Marker, MarkerArray

def rotationx(theta):
    return np.array([
        [1.0, 0.0, 0.0, 0.0],
        [0.0, np.cos(theta / 180 * np.pi), np.sin(theta / 180 * np.pi), 0.0],
        [0.0, -np.sin(theta / 180 * np.pi), np.cos(theta / 180 * np.pi), 0.0],
        [0.0, 0.0, 0.0, 1.0]
    ])

def rotationy(theta):
    return np.array([
        [np.cos(theta / 180 * np.pi), 0.0, np.sin(theta / 180 * np.pi), 0.0],
        [0.0, 1.0, 0.0, 0.0],
        [-np.sin(theta / 180 * np.pi), 0.0, np.cos(theta / 180 * np.pi), 0.0],
        [0.0, 0.0, 0.0, 1.0]
    ])

def rotationz(theta):
    return np.array([
        [np.cos(theta / 180 * np.pi), -np.sin(theta / 180 * np.pi),  0.0, 0.0],
        [np.sin(theta / 180 * np.pi), np.cos(theta / 180 * np.pi), 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0],
        [0.0, 0.0, 0.0, 1.0]
    ])

def draw_mesh(img, intrinsics, model_path, pred_sRT, vertex_color):
    ambient_light = 0.8
    directional_light = 1.0
    scene = pyrender.Scene(ambient_light=np.array([ambient_light, ambient_light, ambient_light, 1.0]), bg_color=[0.0, 0.0, 0.0, 1.0])
    
    direc_l = pyrender.DirectionalLight(color=np.ones(3), intensity=directional_light)
    scene.add(direc_l, pose=np.matmul(rotationy(30), rotationx(45)))
    direc_l = pyrender.DirectionalLight(color=np.ones(3), intensity=directional_light)
    scene.add(direc_l, pose=np.matmul(rotationy(-30), rotationx(45)))
    direc_l = pyrender.DirectionalLight(color=np.ones(3), intensity=directional_light)
    scene.add(direc_l, pose=np.matmul(rotationy(-180), rotationx(45)))
    direc_l = pyrender.DirectionalLight(color=np.ones(3), intensity=(directional_light-0.5))
    scene.add(direc_l, pose=np.matmul(rotationy(0), rotationx(-10)))

    camera = pyrender.IntrinsicsCamera(intrinsics[0][0], intrinsics[1][1], 640 - intrinsics[0][2], 480 - intrinsics[1][2])
    scene.add(camera, pose=np.eye(4))

    for i in range(pred_sRT.shape[0]):
        fuze_trimesh = trimesh.load(model_path)
        fuze_trimesh.vertices = fuze_trimesh.vertices * np.array([[-1.0, -1.0, -1.0]])

        v_max = np.max(fuze_trimesh.vertices, axis=0)
        v_min = np.min(fuze_trimesh.vertices, axis=0)

        color = np.ones([fuze_trimesh.vertices.shape[0], 4]) * 255
        color[:, :3] = (fuze_trimesh.vertices - v_min) * 250 / (v_max - v_min) - 10
        color.astype(np.uint16)
        fuze_trimesh.visual.vertex_colors = color
        # fuze_trimesh.visual.vertex_colors = vertex_color
        # fuze_trimesh.visual.vertex_colors = np.array([100, 100, 100, 255])

        mesh = pyrender.Mesh.from_trimesh(fuze_trimesh)

        # tram = np.eye(3)
        # tram[0,0] = -1
        # tram[1,1] = -1


        R = pred_sRT[i, :3,:3]
        s = np.cbrt(np.linalg.det(pred_sRT[i, :3, :3]))
        R = R / s
        t= pred_sRT[i, :3, 3] / s

        pose=np.eye(4)
        # rotx = rotationx(180)[:3,:3]
        # roty = rotationy(180)[:3,:3]
        # rotz = rotationz(180)[:3,:3]
        # trans = rotx@roty
        pose[:3,:3]= R
        pose[:3,3]= -t

        scene.add(mesh, pose=pose)

    r = pyrender.OffscreenRenderer(img.shape[1], img.shape[0])
    color, _ = r.render(scene)
    color = cv2.flip(color, 1)


    color_mask = np.ones([img.shape[0], img.shape[1], 3], dtype=np.int16)
    color_mask[np.where(color > 0.01)] = 0.0

    img = img * color_mask + color
    img = img.astype(np.uint8)

    return img

def R2q(R:np.array):
    """
    R: 旋转矩阵,左乘
    q: 四元数, q = [x, y, z, w]
    """
        
    w = R[0,0]+R[1,1]+R[2,2]+1
    x = R[0,0]-R[1,1]-R[2,2]+1
    y = -R[0,0]+R[1,1]-R[2,2]+1
    z = -R[0,0]-R[1,1]+R[2,2]+1
    
    q = np.array([w,x,y,z])
    index = np.argmax(q)
    # assert q[index]>0, "max(q) > 0"
    if q[index]<=0:
        return np.array([0,0,0,1],dtype=np.float32)

    q[index] = math.sqrt(q[index]) / 2

    if index==0:
        q0 = q[index]
        q1 = (R[2,1]-R[1,2]) / (4*q0)
        q2 = (R[0,2]-R[2,0]) / (4*q0)
        q3 = (R[1,0]-R[0,1]) / (4*q0)
    elif index==1:
        q1 = q[index]
        q0 = (R[2,1]-R[1,2]) / (4*q1)
        q2 = (R[0,1]+R[1,0]) / (4*q1)
        q3 = (R[2,0]+R[0,2]) / (4*q1)
    elif index==2:
        q2 = q[index]
        q0 = (R[0,2]-R[2,0]) / (4*q2)
        q1 = (R[0,1]+R[1,0]) / (4*q2)
        q3 = (R[1,2]+R[2,1]) / (4*q2)
    elif index==3:
        q3 = q[index]
        q0 = (R[1,0]-R[0,1]) / (4*q3)
        q1 = (R[2,0]+R[0,2]) / (4*q3)
        q2 = (R[1,2]+R[2,1]) / (4*q3)
    else:
        raise ValueError('index error:' +  str(index))
    return np.array([q1,q2,q3,q0],dtype=np.float32)

def q2R(q):
    """
    q: 四元数, q = [x, y, z, w]
    R: 旋转矩阵,左乘
    """
    q = q[[3,0,1,2]]  # (x,y,z,w) -> (w,x,y,z)
    R = np.zeros((3, 3), dtype=float)

    R[0, 0] = 1 - 2 * (q[2] ** 2 + q[3] ** 2)
    R[1, 1] = 1 - 2 * (q[1] ** 2 + q[3] ** 2)
    R[2, 2] = 1 - 2 * (q[1] ** 2 + q[2] ** 2)

    R[0, 1] = 2 * (q[1] * q[2] - q[0] * q[3])
    R[1, 0] = 2 * (q[1] * q[2] + q[0] * q[3])

    R[0, 2] = 2 * (q[1] * q[3] + q[0] * q[2])
    R[2, 0] = 2 * (q[1] * q[3] - q[0] * q[2])

    R[1, 2] = 2 * (q[2] * q[3] - q[0] * q[1])
    R[2, 1] = 2 * (q[2] * q[3] + q[0] * q[1])

    return R

def get_marker_msg(mesh_resource, id, color, scale=1., position=[0,0,0], orientation=[0,0,0,1], frame_name='camera_color_optical_frame'):
    '''
    args:
        position: list or array, [x,y,z]
        orientation: list or array, [x,y,z,w]
        frame_name: str
        scale: float or list/array
        id: used in Marker
        mesh_resource: str
        color: list or array, [r,g,b,a]
    return:
        marker: Marker
    '''
    marker = Marker()
    marker.header.frame_id = frame_name
    marker.pose.position = Point(x=position[0], y=position[1], z=position[2])
    marker.pose.orientation = Quaternion(x=orientation[0], y=orientation[1], z=orientation[2], w=orientation[3])
    marker.type = Marker.MESH_RESOURCE
    marker.id = id
    marker.action = Marker.MODIFY
    marker.mesh_resource = mesh_resource
    if isinstance(scale, float):
        marker.scale.x = scale
        marker.scale.y = scale
        marker.scale.z = scale
    else:
        marker.scale.x = scale[0]
        marker.scale.y = scale[1]
        marker.scale.z = scale[2]
    marker.color.r = color[0]
    marker.color.g = color[1]
    marker.color.b = color[2]
    marker.color.a = color[3]
    return marker

def get_markerarray_msg(marker_list):
    marker_array = MarkerArray()
    marker_array.markers = marker_list
    return marker_array